package com.bryantkobe.myspringcode.utils;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.security.Security;
import java.util.Arrays;
import java.util.Base64;

/**
 * @description: SM4加解密类
 * @author: alienwarezzl
 * @date: 2025/4/24
 */
public class SM4Util {
    private static final String ALGORITHM = "SM4";
    private static final String TRANSFORMATION = "SM4/CBC/PKCS7Padding";
    private static final int KEY_LENGTH = 128; // 128位密钥
    private static final int IV_LENGTH = 16;    // 16字节IV
    private static final int ITERATIONS = 65536; // PBKDF2迭代次数

    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    /**
     * 加密明文
     * @param plaintext 明文
     * @param password  密码
     * @param salt      盐值
     * @return Base64编码的密文（包含IV）
     */
    public static String encrypt(String plaintext, String password, byte[] salt) throws Exception {
        // 生成密钥
        SecretKeySpec secretKey = generateKey(password, salt);

        // 生成随机IV
        byte[] iv = generateIv();
        IvParameterSpec ivSpec = new IvParameterSpec(iv);

        // 加密
        Cipher cipher = Cipher.getInstance(TRANSFORMATION, "BC");
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivSpec);
        byte[] ciphertext = cipher.doFinal(plaintext.getBytes(StandardCharsets.UTF_8));

        // 合并IV和密文
        byte[] combined = new byte[iv.length + ciphertext.length];
        System.arraycopy(iv, 0, combined, 0, iv.length);
        System.arraycopy(ciphertext, 0, combined, iv.length, ciphertext.length);

        return Base64.getEncoder().encodeToString(combined);
    }

    /**
     * 解密密文
     * @param ciphertext Base64编码的密文（包含IV）
     * @param password   密码
     * @param salt       盐值
     * @return 明文
     */
    public static String decrypt(String ciphertext, String password, byte[] salt) throws Exception {
        // 解码Base64
        byte[] combined = Base64.getDecoder().decode(ciphertext);

        // 分离IV和密文
        byte[] iv = Arrays.copyOfRange(combined, 0, IV_LENGTH);
        byte[] ct = Arrays.copyOfRange(combined, IV_LENGTH, combined.length);

        // 生成密钥
        SecretKeySpec secretKey = generateKey(password, salt);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);

        // 解密
        Cipher cipher = Cipher.getInstance(TRANSFORMATION, "BC");
        cipher.init(Cipher.DECRYPT_MODE, secretKey, ivSpec);
        byte[] plaintext = cipher.doFinal(ct);

        return new String(plaintext, StandardCharsets.UTF_8);
    }

    /**
     * 使用PBKDF2生成SM4密钥
     */
    private static SecretKeySpec generateKey(String password, byte[] salt) throws Exception {
        PBEKeySpec spec = new PBEKeySpec(password.toCharArray(), salt, ITERATIONS, KEY_LENGTH);
        SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256");
        byte[] keyBytes = factory.generateSecret(spec).getEncoded();
        return new SecretKeySpec(keyBytes, ALGORITHM);
    }

    /**
     * 生成随机IV
     */
    private static byte[] generateIv() {
        byte[] iv = new byte[IV_LENGTH];
        new SecureRandom().nextBytes(iv);
        return iv;
    }

    // 示例用法
    public static void main(String[] args) throws Exception {
        String plainText = "Hello, SM4!";
        String password = "mySecretPassword";
        byte[] salt = "customSalt123".getBytes(StandardCharsets.UTF_8); // 自定义盐值

        // 加密
        String encrypted = encrypt(plainText, password, salt);
        System.out.println("加密结果: " + encrypted);

        // 解密
        String decrypted = decrypt(encrypted, password, salt);
        System.out.println("解密结果: " + decrypted);
    }
}
